import os
import io
from io import BytesIO
import uuid
import PIL.Image
import codecs
import threadpool
import numpy as np
import cv2
import base64
from math import radians, fabs, sin, cos
from labelme.label_file import *
from labelme import utils
import shutil

dirpath = "phone_detect/labelme_annotations"
dst_path_dir = "phone_detect/yolo_annotations"

files = [os.path.join(y, file) for y, z, x in os.walk(dirpath)
         for file in x if os.path.splitext(file)[1] == '.json']

if not os.path.exists(dst_path_dir):
    os.makedirs(dst_path_dir)

index_files = []
for i in range(len(files)):
    index_files.append([i, files[i]])


def ThreadFun_c1(file):
    index = file[0]
    file = file[1]
    print(file)
    txtname = os.path.join(dst_path_dir, '{}.txt'.format(index))
    label_file = LabelFile()
    label_file.load(file)

    arr = utils.img_b64_to_arr(base64.b64encode(label_file.imageData).decode("utf-8"))
    label_file.imageHeight, label_file.imageWidth, _ = arr.shape
    shutil.copy(file.replace('.json', '.jpg'), txtname.replace('.txt', '.jpg'))
    with codecs.open(txtname, 'w', 'GBK') as f:
        for shape in label_file.shapes:
            x1 = shape['points'][0][0] / label_file.imageWidth
            y1 = shape['points'][0][1] / label_file.imageHeight
            x2 = shape['points'][1][0] / label_file.imageWidth
            y2 = shape['points'][1][1] / label_file.imageHeight
            f.write('{} {} {} {} {}\n'.format(shape['label'], (x1 + x2) / 2, (y1 + y2) / 2, abs(x1 - x2), abs(y1 - y2)))


# 定义了一个线程池，最多创建8个线程
pool = threadpool.ThreadPool(16)
# 创建要开启多线程的函数，以及函数相关参数和回调函数，其中回调数可以不写，default是none
requests = threadpool.makeRequests(ThreadFun_c1, index_files)
# 将所有要运行多线程的请求扔进线程池
[pool.putRequest(req) for req in requests]
# 所有的线程完成工作后退出
pool.wait()

'''
检查标注文件对不对
'''
for file in index_files:
    index = file[0]
    file = file[1]
    txtname = os.path.join(dst_path_dir, '{}.txt'.format(index))
    imgfile = txtname.replace('.txt', '.jpg')
    img = cv2.imread(imgfile)
    h, w = img.shape[:2]
    with codecs.open(txtname, 'r', 'GBK') as f:
        liness = f.readlines()
        for line in liness:
            lines = line.strip('\n').split(' ')
            for i in range(1, 5):
                lines[i] = float(lines[i])
            x1 = (lines[1] - lines[3] / 2) * w
            x2 = (lines[1] + lines[3] / 2) * w
            y1 = (lines[2] - lines[4] / 2) * h
            y2 = (lines[2] + lines[4] / 2) * h

            x1 = int(x1)
            x2 = int(x2)
            y1 = int(y1)
            y2 = int(y2)

            if lines[0] == '0':
                color = (0, 255, 0)
            else:
                color = (0, 0, 255)

            cv2.rectangle(img, (x1, y1), (x2, y2), color)
    cv2.imshow('1', img)
    cv2.waitKey(0)
